import numpy as np
from sklearn.cluster import KMeans
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import SVC


#训练模型  预测模型

class ModelManagement:

    def __init__(self):
        pass


    @staticmethod
    def trainClassificationModel(data, tag, module_name, proportion):
        dataArray = np.array(data)
        tagArray = np.array(tag)
        X_train, X_test, y_train, y_test = train_test_split(dataArray, tagArray, test_size=proportion)
        if module_name == "Naive Bayes":
            gnb = GaussianNB()
            gnb.fit(X_train, y_train)
            return gnb
        elif module_name == "Kmeans":
            km = KMeans(n_clusters=2)
            km.fit(X_train)
            return km
        elif module_name == "SVM":
            clf_0 = SVC(kernel='rbf', random_state=0, gamma=1, C=1)
        # clf_0 = SVC(C=1,cache_size=200,class_weight=None,coef0=0.0,
        #     decision_function_shape='ovr',degree=3,gamma='auto',kernel='linear',
        #     max_iter=-1,probability=False,random_state=0,shrinking=True,
        #     tol=0.001,verbose=False)
        # clf_1 = SVC(C=1, kernel='rbf', gamma=20, decision_function_shape='ovr') # 使用rbf径向基函数来讲低维数据转化为高维数据，使其可分
        #clf_2 = SVC(C=1, kernel='linear', gamma=20, decision_function_shape='ovr')
        #clf_3 =  SVC(C=1, kernel='poly', gamma=20, decision_function_shape='ovr')
        # clf_4 = SVC(C=1, kernel='sigmoid', gamma=20, decision_function_shape='ovr')
         #   print(y_train)
        #    print(X_train)
            clf_0.fit(X_train, y_train)
            return clf_0
        else:
            print("not module you change")
            return



    def modelPredict(self, predict_data, train_data, tag, config, if_save):
        model = self.trainClassificationModel(train_data,tag, config.get("model_name"),config.get("proportion"))
        result = model.predict(predict_data)
       # print(result)
        if if_save:
            F = open(r'20201217022021.txt', 'w')
            for i in result:
                F.write(str(i)+",")
            F.close()
        return result



